import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import gradio as gr

from utils import zip_folder
from model import ModelFactory
from translator import PDFTranslator

def translate_pdf(model_type, model, target_language, api_key, filepath):
    model = ModelFactory.create_model(model_type, model, api_key)
    output_file_path = os.path.splitext(filepath)[0]
    translator = PDFTranslator(model)
    translator.translate_pdf(filepath, 'markdown', output_file_path, target_language)
    zip_file_path = zip_folder(output_file_path, os.path.dirname(filepath))
    if zip_file_path:
        return zip_file_path
    else:
        gr.Error('Failed to create zip file')


if __name__ == "__main__":
    translator_gui = gr.Interface(
        fn=translate_pdf,
        title='Mini AI Translator',
        inputs=[
            gr.Dropdown(
                ['OpenAIModel', 'ChatGLMModel'], label='model_type',
                info='The type of translation model to use. Choose between "OpenAIModel" and "ChatGLMModel".'
            ),
            gr.Dropdown(
                ['gpt-3.5-turbo', 'glm-4'], label='model',
                info='The specific model to use for translation. For example, "gpt-3.5-turbo" for OpenAIModel or "glm-4" for ChatGLMModel.'
            ),
            gr.Dropdown(
                ['Chinese', 'Japanese'], label='target_language', info='Choose the target language for translation.'
            ),
            gr.Textbox(label="api_key", info='API key corresponding to the selected model_type.'),
            gr.File(label='Upload the PDF file that needs to be translated'),
        ],
        outputs=[
            gr.File(label='Download the translated file zip')
        ],
        examples=[
            ['OpenAIModel', 'gpt-3.5-turbo', 'Chinese', 'xxx', None],
            ['ChatGLMModel', 'glm-4', 'Japanese', 'xxx', None],
        ]
    )

    translator_gui.launch(share=True)
